package com.wondersgroup.aida.embeddingdemo;

import com.wondersgroup.aida.tools.ModelUtil;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

/**
 * @author: 紫麒麟vip
 * @create: 2025/2/16
 * Description:
 */
public class QuestionSearchDemo {
    public static void main(String[] args) {
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ModelUtil.BASE_URI_OPENAI)
                .apiKey(ModelUtil.API_KEY_0PENAI)
                .build();

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .build();
        //预设几个指示，生成向量
        TextSegment textSegment1 = TextSegment.textSegment("客服电话是400-8888888");
        TextSegment textSegment2 = TextSegment.textSegment("客服工作时间是周一到周五");
        TextSegment textSegment3 = TextSegment.textSegment("客服投诉电话是400-6666666");
        Response<Embedding> embed1 = embeddingModel.embed(textSegment1);
        Response<Embedding> embed2 = embeddingModel.embed(textSegment2);
        Response<Embedding> embed3 = embeddingModel.embed(textSegment3);
        //存储向量
        embeddingStore.add(embed1.content(), textSegment1);
        embeddingStore.add(embed2.content(), textSegment2);
        embeddingStore.add(embed3.content(), textSegment3);
        //预设一个问题，生成向量
        Response<Embedding> embed = embeddingModel.embed("客服电话多少");
        //查询
        List<EmbeddingMatch<TextSegment>> result = embeddingStore.findRelevant(embed.content(), 10,5);
        for (EmbeddingMatch<TextSegment> embeddingMatch : result) {
            System.out.println(embeddingMatch.embedded().text() + "，分数为：" + embeddingMatch.score());
        }
    }
}
